{
 "cells": [
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.applications import MobileNetV2\n",
    "from tensorflow.keras.layers import Dense, GlobalAveragePooling2D\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "# 设置随机种子，保证结果可复现\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# 数据集路径（这里使用TensorFlow的flowers数据集）\n",
    "dataset_url = \"https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz\"\n",
    "data_dir = tf.keras.utils.get_file('flower_photos', origin=dataset_url, untar=True)\n",
    "data_dir = os.path.join(data_dir)\n",
    "\n",
    "# 超参数设置\n",
    "IMAGE_SIZE = (224, 224)  # MobileNetV2默认输入尺寸\n",
    "BATCH_SIZE = 32\n",
    "EPOCHS = 20\n",
    "VALIDATION_SPLIT = 0.2\n",
    "\n",
    "# 数据增强 - 防止过拟合\n",
    "train_datagen = ImageDataGenerator(\n",
    "    rescale=1./255,\n",
    "    validation_split=VALIDATION_SPLIT,\n",
    "    rotation_range=20,\n",
    "    width_shift_range=0.2,\n",
    "    height_shift_range=0.2,\n",
    "    horizontal_flip=True,\n",
    "    zoom_range=0.2\n",
    ")\n",
    "\n",
    "# 加载训练集\n",
    "train_generator = train_datagen.flow_from_directory(\n",
    "    data_dir,\n",
    "    target_size=IMAGE_SIZE,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    class_mode='categorical',\n",
    "    subset='training'\n",
    ")\n",
    "\n",
    "# 加载验证集\n",
    "validation_generator = train_datagen.flow_from_directory(\n",
    "    data_dir,\n",
    "    target_size=IMAGE_SIZE,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    class_mode='categorical',\n",
    "    subset='validation'\n",
    ")\n",
    "\n",
    "# 获取类别名称\n",
    "class_names = list(train_generator.class_indices.keys())\n",
    "num_classes = len(class_names)\n",
    "print(f\"类别名称: {class_names}\")\n",
    "print(f\"类别数量: {num_classes}\")\n",
    "\n",
    "# 加载预训练的MobileNetV2模型（不包含顶层分类器）\n",
    "base_model = MobileNetV2(\n",
    "    input_shape=(*IMAGE_SIZE, 3),\n",
    "    include_top=False,  # 不包含顶层分类器\n",
    "    weights='imagenet'  # 使用ImageNet上预训练的权重\n",
    ")\n",
    "\n",
    "# 冻结基础模型的权重（迁移学习）\n",
    "base_model.trainable = False\n",
    "\n",
    "# 添加自定义分类层\n",
    "x = base_model.output\n",
    "x = GlobalAveragePooling2D()(x)  # 全局平均池化\n",
    "x = Dense(128, activation='relu')(x)  # 全连接层\n",
    "predictions = Dense(num_classes, activation='softmax')(x)  # 输出层\n",
    "\n",
    "# 构建完整模型\n",
    "model = Model(inputs=base_model.input, outputs=predictions)\n",
    "\n",
    "# 编译模型\n",
    "model.compile(\n",
    "    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),\n",
    "    loss='categorical_crossentropy',\n",
    "    metrics=['accuracy']\n",
    ")\n",
    "\n",
    "# 查看模型结构\n",
    "model.summary()\n",
    "\n",
    "# 训练模型\n",
    "history = model.fit(\n",
    "    train_generator,\n",
    "    steps_per_epoch=train_generator.samples // BATCH_SIZE,\n",
    "    validation_data=validation_generator,\n",
    "    validation_steps=validation_generator.samples // BATCH_SIZE,\n",
    "    epochs=EPOCHS\n",
    ")\n",
    "\n",
    "# 微调：解冻部分顶层进行训练（可选）\n",
    "# base_model.trainable = True\n",
    "# # 只解冻顶层的部分层\n",
    "# for layer in base_model.layers[:-20]:\n",
    "#     layer.trainable = False\n",
    "#\n",
    "# # 使用较小的学习率进行微调\n",
    "# model.compile(\n",
    "#     optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),\n",
    "#     loss='categorical_crossentropy',\n",
    "#     metrics=['accuracy']\n",
    "# )\n",
    "#\n",
    "# # 继续训练\n",
    "# history_fine = model.fit(\n",
    "#     train_generator,\n",
    "#     steps_per_epoch=train_generator.samples // BATCH_SIZE,\n",
    "#     validation_data=validation_generator,\n",
    "#     validation_steps=validation_generator.samples // BATCH_SIZE,\n",
    "#     epochs=EPOCHS + 10\n",
    "# )\n",
    "\n",
    "# 绘制训练过程中的准确率和损失曲线\n",
    "def plot_training_history(history):\n",
    "    acc = history.history['accuracy']\n",
    "    val_acc = history.history['val_accuracy']\n",
    "    loss = history.history['loss']\n",
    "    val_loss = history.history['val_loss']\n",
    "\n",
    "    epochs_range = range(len(acc))\n",
    "\n",
    "    plt.figure(figsize=(12, 4))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(epochs_range, acc, label='Training Accuracy')\n",
    "    plt.plot(epochs_range, val_acc, label='Validation Accuracy')\n",
    "    plt.legend(loc='lower right')\n",
    "    plt.title('Training and Validation Accuracy')\n",
    "\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(epochs_range, loss, label='Training Loss')\n",
    "    plt.plot(epochs_range, val_loss, label='Validation Loss')\n",
    "    plt.legend(loc='upper right')\n",
    "    plt.title('Training and Validation Loss')\n",
    "    plt.savefig('training_history.png')\n",
    "    plt.show()\n",
    "\n",
    "# 绘制训练历史\n",
    "plot_training_history(history)\n",
    "\n",
    "# 保存模型\n",
    "model.save('flower_classifier_mobilenetv2.h5')\n",
    "print(\"模型已保存为 'flower_classifier_mobilenetv2.h5'\")\n",
    "\n",
    "# 评估模型在验证集上的表现\n",
    "loss, accuracy = model.evaluate(validation_generator)\n",
    "print(f\"验证集损失: {loss:.4f}\")\n",
    "print(f\"验证集准确率: {accuracy:.4f}\")\n"
   ],
   "id": "930cdfe01411fde5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
